import numpy as np
import math
from scipy.spatial.transform import Rotation

def quaternion_to_rotation_matrix(quaternion, order='wxyz'):
    # calculate rotation matrix from quaternion
    if order == 'wxyz':
        w, x, y, z = quaternion
    elif order == 'xyzw':
        x, y, z, w = quaternion
    else:
        return

    R = np.array([
        [1 - 2*y**2 - 2*z**2,     2*x*y - 2*z*w,     2*x*z + 2*y*w],
        [2*x*y + 2*z*w,     1 - 2*x**2 - 2*z**2,     2*y*z - 2*x*w],
        [2*x*z - 2*y*w,     2*y*z + 2*x*w,     1 - 2*x**2 - 2*y**2]
    ])

    return R


def euler_to_rotation_matrix(euler, is_degree=False):
    # calculate rotation matrix from euler angles
    roll, pitch, yaw = euler
    if is_degree:
        roll = np.deg2rad(roll)
        pitch = np.deg2rad(pitch)
        yaw = np.deg2rad(yaw)

    R_x = np.array([
        [1, 0, 0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll), np.cos(roll)],
    ])
    R_y = np.array([
        [np.cos(pitch), 0, np.sin(pitch)],
        [0, 1, 0],
        [-np.sin(pitch), 0, np.cos(pitch)],
    ])
    R_z = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw), np.cos(yaw), 0],
        [0, 0, 1],
    ])

    R = R_z @ R_y @ R_x

    return R


def isRotationMatrix(R):
    # Checks if a matrix is a valid rotation matrix.
    Rt = np.transpose(R)
    shouldBeIdentity = np.dot(Rt, R)
    I = np.identity(3, dtype=R.dtype)
    n = np.linalg.norm(I - shouldBeIdentity)

    return n < 1e-6


def rotation_matrix_to_euler(R):
    # calculate euler angles from rotation matrix
    assert isRotationMatrix(R)
    sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = math.atan2(R[2, 1], R[2, 2])
        y = math.atan2(-R[2, 0], sy)
        z = math.atan2(R[1, 0], R[0, 0])
    else:
        x = math.atan2(-R[1, 2], R[1, 1])
        y = math.atan2(-R[2, 0], sy)
        z = 0

    return np.array([x, y, z], dtype=float)

def rotation_matrix_to_quaternion(R):
    # calculate quaternion from rotation matrix
    tr = R[0, 0] + R[1, 1] + R[2, 2]

    if tr > 0:
        w = np.sqrt(1.0 + tr) / 2.0
        x = (R[2, 1] - R[1, 2]) / (4.0 * w)
        y = (R[0, 2] - R[2, 0]) / (4.0 * w)
        z = (R[1, 0] - R[0, 1]) / (4.0 * w)
    else:
        # Find the largest diagonal element
        ii = np.argmax([R[0, 0], R[1, 1], R[2, 2]])

        if ii == 0:
            w = (R[0, 0] - R[1, 1] - R[2, 2] + 1.0) / 4.0
            x = np.sqrt(w)
            y = (R[0, 1] + R[1, 0]) / (4.0 * x)
            z = (R[0, 2] + R[2, 0]) / (4.0 * x)
        elif ii == 1:
            w = (R[1, 1] - R[0, 0] - R[2, 2] + 1.0) / 4.0
            x = (R[0, 1] + R[1, 0]) / (4.0 * w)
            y = np.sqrt(w)
            z = (R[1, 2] + R[2, 1]) / (4.0 * w)
        elif ii == 2:
            w = (R[2, 2] - R[0, 0] - R[1, 1] + 1.0) / 4.0
            x = (R[0, 2] + R[2, 0]) / (4.0 * w)
            y = (R[1, 2] + R[2, 1]) / (4.0 * w)
            z = np.sqrt(w)

    return np.array([x, y, z, w], dtype=float)


def state2SE3(states, is_quaternion=False):
    # states = [x, y, z, qx, qy, qz, qw] or [x, y, z, roll, pitch, yaw]
    if is_quaternion:
        x, y, z, qx, qy, qz, qw = states
        rotation_matrix = quaternion_to_rotation_matrix([qw, qx, qy, qz])
    else:
        x, y, z, roll, pitch, yaw = states
        rotation_matrix = euler_to_rotation_matrix([roll, pitch, yaw])

    transformation = np.diag([1.0, 1.0, 1.0, 1.0])
    transformation[:3, :3] = rotation_matrix
    transformation[:3, 3] = [x, y, z]
    
    return transformation

def sE32states(transformation, to_quaternion=False):
    # states = [x, y, z, qx, qy, qz, qw] or [x, y, z, roll, pitch, yaw]
    if to_quaternion:
        states = np.zeros(7, dtype=float)
        states[3:] = Rotation.from_matrix(transformation[:3, :3]).as_quat()
    else:
        states = np.zeros(6, dtype=float)
        states[3:] = Rotation.from_matrix(transformation[:3, :3]).as_euler(
            "xyz", degrees=False
        )
    
    states[:3] = transformation[:3, 3]

    return states
